import os
import sys
import torch
from PIL import Image
from flask import Flask, request, render_template
from flask_cors import CORS
from werkzeug.utils import secure_filename
import webbrowser
import time

dir_mytest = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, dir_mytest)

# from web_predict import init_model, predict
from detect_web import webber

app = Flask(__name__)
CORS(app)  # 解决跨域问题
pa = './'

upload_path = 'F:\WIT\YY\yolov5\web\static\imgs'
generate_path = 'F:\WIT\YY\yolov5\web\static\gen'

data_name = ''

if not os.path.exists(upload_path):
    os.mkdir(upload_path)

if not os.path.exists(generate_path):
    os.mkdir(generate_path)


@app.route("/predict", methods=["POST"])
@torch.no_grad()
def solve():
    global data_name
    image = request.files["file"]
    data_name = secure_filename(image.filename)
    up_path = os.path.join(upload_path, data_name)
    gen_path = os.path.join(generate_path, data_name)
    image.save(up_path)

    res_dir = webber(up_path)
    res = os.path.join(res_dir,data_name)

    new_img = Image.open(res)
    gen_path = os.path.join(generate_path,data_name)
    # # os.path.join()

    # TODO 
    # new_img.show()
    new_img.save(gen_path)
    return {'data_name': data_name}


@app.route("/", methods=["GET", "POST"])
def root():
    # TODO
    return render_template("up.html",data_name = data_name)
    # return render_template("up.html",path = gen_path)



if __name__ == '__main__':
    port = 5000
    webbrowser.open_new(f'http://127.0.0.1:{port}/')
    app.run(host="0.0.0.0", port=port)



